package com.nan.bytecode.asm;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.commons.AdviceAdapter;

import java.io.IOException;

public class AOPImprove {
    public static byte[] addProfiling() throws IOException {
        ClassReader classReader = new ClassReader(UserModel.class.getName());
        ClassWriter classWriter = new ClassWriter(classReader, ClassWriter.COMPUTE_MAXS);
        ClassVisitor classVisitor = new ProfilingClassVisitor(classWriter);
        classReader.accept(classVisitor, ClassReader.EXPAND_FRAMES);
        return classWriter.toByteArray();
    }

    static class ProfilingClassVisitor extends ClassVisitor {
        public ProfilingClassVisitor(ClassVisitor classVisitor) {
            super(Opcodes.ASM7, classVisitor);
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            System.out.println("access=" + access + ", name=" + name + ", descriptor=" + descriptor);
            MethodVisitor methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions);
            if (!"queryUserInfo".equals(name)) return methodVisitor;
            return new ProfilingMethodVisitor(methodVisitor, access, name, descriptor);
        }
    }

    static class ProfilingMethodVisitor extends AdviceAdapter {
        protected ProfilingMethodVisitor(MethodVisitor methodVisitor, int access, String name, String descriptor) {
            super(ASM7, methodVisitor, access, name, descriptor);
        }

        @Override
        protected void onMethodEnter() {
            // System.nanoTime();
            mv.visitMethodInsn(Opcodes.INVOKESTATIC, "java/lang/System", "nanoTime", "()J", false);
            // 存储在索引为2的局部变量表，注意：对应方法中只能包含一个参数，且不能包含其他局部变量，否则其他局部变量有可能已占用了该索引，从而引起问题
            // 对于直接设置较大的索引，也是不允许的，必须顺序存储。
            mv.visitVarInsn(Opcodes.LSTORE, 2);
        }

        @Override
        protected void onMethodExit(int opcode) {
            // 如果是返回语句或者抛出异常语句，则表示方法执行完成
            System.out.println("opcode=" + opcode);
            if (opcode >= Opcodes.IRETURN && opcode <= Opcodes.RETURN || opcode == Opcodes.ATHROW) {
                mv.visitFieldInsn(Opcodes.GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
                mv.visitMethodInsn(Opcodes.INVOKESTATIC, "java/lang/System", "nanoTime", "()J", false);
                mv.visitVarInsn(Opcodes.LLOAD, 2);
                mv.visitInsn(Opcodes.LSUB);
                mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(J)V", false);
            }
        }
    }
}
